{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Implementing an AI Agent with SmolAgents\n",
        "This notebook is a companion of chapter 12 of the \"Small Domain Specific LLMs in Action\" book, author Guglielmo Iozzia, [Manning Publications](https://www.manning.com/), 2025.  \n",
        "The code in this notebook is about implementing a basic AI Agent using the Hugging Face's [SmolAgents](https://github.com/huggingface/smolagents) framework and Small Language Models (SLMs). Execution of the code cells in this notebook requires hardware acceleration.   \n",
        "More details about the code can be found in the related book's chapter."
      ],
      "metadata": {
        "id": "B0v4TyCAs3-T"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Install the missing requirements (only SmolAgents missing in the Colab VM)."
      ],
      "metadata": {
        "id": "gTxCLXO_uucs"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "epI7BfDag4fJ"
      },
      "outputs": [],
      "source": [
        "!pip install smolagents"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Import the required packages and classes."
      ],
      "metadata": {
        "id": "cOoeTT6svGO1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "from smolagents import CodeAgent, TransformersModel, tool"
      ],
      "metadata": {
        "id": "EOjno8KpiURY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Because this notebook is for learning purposes only, let's define a mock webpage containing a list of flights, to prevent the agent running longer queries across the web."
      ],
      "metadata": {
        "id": "knpHAxxgvJeB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "mock_flights_web_page = \"\"\"\n",
        "1. Flight DUB-BOS123\n",
        "- Price: $350\n",
        "- Origin: Dublin Airport (DUB)\n",
        "- Destination: Boston Logan International Airport (BOS)\n",
        "- Date: January 10, 2025\n",
        "\n",
        "2. Flight DUB-EWR456\n",
        "- Price: $370\n",
        "- Origin: Dublin Airport (DUB)\n",
        "- Destination: Newark Liberty International Airport (EWR)\n",
        "- Date: January 10, 2025\n",
        "\n",
        "3. Flight DUB-JFK789\n",
        "- Price: $400\n",
        "- Origin: Dublin Airport (DUB)\n",
        "- Destination: John F. Kennedy International Airport (JFK)\n",
        "- Date: January 20, 2025\n",
        "\n",
        "4. Flight DUB-JFK101\n",
        "- Price: $450\n",
        "- Origin: Dublin Airport (DUB)\n",
        "- Destination: John F. Kennedy International Airport (JFK)\n",
        "- Date: January 10, 2025\n",
        "\n",
        "5. Flight JFK-DUB202\n",
        "- Price: $200\n",
        "- Origin: John F. Kennedy International Airport (JFK)\n",
        "- Destination: Dublin Airport (DUB)\n",
        "- Date: January 12, 2025\n",
        "\n",
        "6. Flight GWY-PHL303\n",
        "- Price: $400\n",
        "- Origin: Galway Airport (GWY)\n",
        "- Destination: Philadelphia International Airport (PHL)\n",
        "- Date: January 12, 2025\n",
        "\n",
        "7. Flight GWY-CDG404\n",
        "- Price: $150\n",
        "- Origin: Galway Airport (GWY)\n",
        "- Destination: Charles de Gaulle International Airport (CDG)\n",
        "- Date: January 10, 2025\n",
        "\n",
        "8. Flight DUB-PHL505\n",
        "- Price: $390\n",
        "- Origin: Dublin Airport (DUB)\n",
        "- Destination: Philadelphia International Airport (PHL)\n",
        "- Date: January 10, 2025\n",
        "\"\"\""
      ],
      "metadata": {
        "id": "YGGi1WyWVLyS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's implement the first custom tool for the agent. We achieve this by defining a function, to be decorated with the `@tool` decorator, that can extract from a web page a list of flights that match a given origin and destination pair on a given date. Please note that SmolAgents, for any tool implemented this way, requires type hints for inputs and outputs and docstrings that explain what the function does and what the arguments and return values are."
      ],
      "metadata": {
        "id": "Ls4nK-ltveob"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import re\n",
        "\n",
        "@tool\n",
        "def search_for_flights(origin: str, destination: str, date: str, web_page: str) -> str:\n",
        "  \"\"\"Searches for flights based on origin, destination, and date.\n",
        "\n",
        "  Args:\n",
        "    origin: The origin airport code.\n",
        "    destination: The destination airport code.\n",
        "    date: The travel date.\n",
        "    web_page: The webpage content to search for flights.\n",
        "\n",
        "  Returns:\n",
        "    The flight code of the matching flight.\n",
        "  \"\"\"\n",
        "\n",
        "  origin = origin.lower()\n",
        "  destination = destination.lower()\n",
        "\n",
        "  web_page = mock_flights_web_page\n",
        "  flights = web_page.strip().split('\\n\\n')\n",
        "  matching_flights = []\n",
        "\n",
        "  matching_pattern = r\"Flight\\s([A-Z]*-[A-Z]*[0-9]*)- Price:\\s+\\$(\\d+)- Origin:\\s+(.*?)- Destination:\\s+(.*?)- Date:\\s+([a-zA-Z]+\\s[0-9]+,\\s[0-9]+)\"\n",
        "  flight_pattern = re.compile(\n",
        "      matching_pattern\n",
        "      )\n",
        "\n",
        "  for flight in flights:\n",
        "      flight = flight.replace('\\n', '')\n",
        "      match = flight_pattern.search(flight)\n",
        "      if match:\n",
        "          flight_code = match.group(1)\n",
        "          price = match.group(2)\n",
        "          flight_origin = match.group(3).lower()\n",
        "          flight_destination = match.group(4).lower()\n",
        "          flight_date = match.group(5)\n",
        "\n",
        "          if origin in flight_origin and destination in flight_destination:\n",
        "              matching_flights.append({\n",
        "                  \"Flight Code\": flight_code,\n",
        "                  \"Price\": price,\n",
        "                  \"Origin\": flight_origin,\n",
        "                  \"Destination\": flight_destination,\n",
        "                  \"Date\": flight_date\n",
        "              })\n",
        "\n",
        "  return matching_flights"
      ],
      "metadata": {
        "id": "DU7H8eHxTIFe"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The code cell below is to test the tool to search for flights. Execution isn't mandatory to run the agent."
      ],
      "metadata": {
        "id": "q7v0niFlv1b2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "matching_flight = search_for_flights('DUB', 'JFK', '2023-03-15', mock_flights_web_page)\n",
        "matching_flight"
      ],
      "metadata": {
        "id": "TpgIAAWeCbb3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define a function (tool) to find a seat within a given flight and class (Economy or Business)."
      ],
      "metadata": {
        "id": "AFyKsxkdwHO6"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import random\n",
        "\n",
        "@tool\n",
        "def find_seat(flight: str, seat_class: str='economy') -> str:\n",
        "  \"\"\"Finds a seat within a given flight.\n",
        "\n",
        "  Args:\n",
        "    flight: The flight code.\n",
        "    seat_class: The desired seat class.\n",
        "\n",
        "  Returns:\n",
        "    The seat number.\n",
        "  \"\"\"\n",
        "\n",
        "  if seat_class.lower() not in ['business', 'economy']:\n",
        "        raise ValueError(\"ticket_type must be either 'business' or 'economy'\")\n",
        "\n",
        "  seat_type = random.choice(['A', 'B', 'C', 'D', 'E', 'F'])\n",
        "\n",
        "  if seat_class.lower() == 'business':\n",
        "        seat_number = random.randint(1, 3)\n",
        "  else:\n",
        "      seat_number = random.randint(4, 32)\n",
        "\n",
        "  selected_seat = f\"{seat_type}{seat_number}\"\n",
        "\n",
        "  return selected_seat\n"
      ],
      "metadata": {
        "id": "k3NrfyH2UV1y"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The code cell below is to test the tool to find a seat for a given flight. Execution isn't mandatory to run the agent."
      ],
      "metadata": {
        "id": "9cCnhkYCwWMU"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "selected_seat = find_seat('SFO-ANC123', 'Economy')\n",
        "selected_seat"
      ],
      "metadata": {
        "id": "EXJu3VlWad55"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define a function (tool) to find book a ticket for a given flight and seat."
      ],
      "metadata": {
        "id": "h4Xlj8i_we2h"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "@tool\n",
        "def book_flight_ticket(flight: str, seat: str) -> str:\n",
        "  \"\"\"Books a flight ticket.\n",
        "\n",
        "  Args:\n",
        "    flight: The flight code.\n",
        "    seat: The seat number.\n",
        "\n",
        "  Returns:\n",
        "    The ticket number.\n",
        "  \"\"\"\n",
        "  ticket_number = random.randint(100000, 999999)\n",
        "\n",
        "  return f\"Ticket booked for flight {flight} and seat {seat}. Ticket number: {ticket_number}\""
      ],
      "metadata": {
        "id": "AVfv07PsxL67"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The code cell below is to test the tool to book a flight ticket. Execution isn't mandatory to run the agent."
      ],
      "metadata": {
        "id": "a0AKKY1jxCJG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "booking = book_flight_ticket('DUB-JFK789', 'A1')\n",
        "booking"
      ],
      "metadata": {
        "id": "VcDdmtjyxqbv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define a function (tool) to calculate the n-th Fibonacci number."
      ],
      "metadata": {
        "id": "mQpvyQcSxFt3"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "@tool\n",
        "def fibonacci(n: int) -> int:\n",
        "    \"\"\"\n",
        "    Calculates the n-th Fibonacci number.\n",
        "\n",
        "    Args:\n",
        "      n: The sequence starting number.\n",
        "\n",
        "    Returns:\n",
        "      The n-th number in the sequence.\n",
        "    \"\"\"\n",
        "    if n < 0:\n",
        "        raise ValueError(\"Negative arguments are not supported\")\n",
        "    elif n == 0:\n",
        "        return 0\n",
        "    elif n == 1:\n",
        "        return 1\n",
        "\n",
        "    a, b = 0, 1\n",
        "    for _ in range(2, n + 1):\n",
        "        a, b = b, a + b\n",
        "    return b"
      ],
      "metadata": {
        "id": "FSftdp6VdNKp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The code cell below is to test the tool to calculate the n-th Fibonacci number. Execution isn't mandatory to run the agent."
      ],
      "metadata": {
        "id": "pLNlR5R2xRZf"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "n = 10\n",
        "print(f\"The {n}-th Fibonacci number is: {fibonacci(n)}\")"
      ],
      "metadata": {
        "id": "z4Pgxq8tdZZT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Download the instructed model ([HuggingFaceTB/SmolLM2-1.7B-Instruct](https://huggingface.co/HuggingFaceTB/SmolLM2-1.7B-Instruct)) to be used by the Agent."
      ],
      "metadata": {
        "id": "d69ND38ZxaQH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "model = TransformersModel(\"HuggingFaceTB/SmolLM2-1.7B-Instruct\",\n",
        "                          device_map=\"auto\",\n",
        "                          max_new_tokens=1000,\n",
        "                          torch_dtype=torch.float16)"
      ],
      "metadata": {
        "id": "MDTomdqaiXUj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Set up a CodeAgent, adding to its toolset the three custom tools for managing a flight search and booking and specifying the SLM to use."
      ],
      "metadata": {
        "id": "d-IX_TYlx0AU"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "custom_tools = [search_for_flights, find_seat, book_flight_ticket]\n",
        "agent = CodeAgent(tools=custom_tools, model=model,\n",
        "                  verbosity_level=2, max_steps=3)"
      ],
      "metadata": {
        "id": "Ov_Sp2Rke0Dd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Run the agent on a user task and observe how it analyze it, splits it into steps, selects the tools, generates Python code and executes it."
      ],
      "metadata": {
        "id": "WV-CznJByGRH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "user_task = \"Find a flight from Dublin to New York and then select a seat.\"\n",
        "agent.run(user_task)"
      ],
      "metadata": {
        "id": "8xm_Q8D5iacH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "4FYW9SfJybhe"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}